package com.example.nplboot.cs;

import cn.hutool.core.io.FileUtil;
import com.aliasi.classify.*;
import com.example.nplboot.model.FileModel;
import com.hankcs.hanlp.HanLP;
import com.hankcs.hanlp.seg.common.Term;
import org.apache.poi.hwpf.extractor.WordExtractor;
import org.apache.poi.xwpf.extractor.XWPFWordExtractor;
import org.apache.poi.xwpf.usermodel.XWPFDocument;

import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.nio.charset.StandardCharsets;
import java.text.NumberFormat;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 * 测试运行方法，直接运行
 * @Classname CSTest
 * @Date 2023/2/12 19:49
 * @Created by 87766867@qq.com
 */
public class CSTest {

    private  static String[] CATEGORIES = {
            "原告被告",
            "原告被告代理人",
            "原告被告代理人代表人",
            "原告被告代表人"
    };


    //模型目录
    private static String MODEL_FILE = "/数据目录/训练好的模型/tfidf_classifier";

    //测试文件目录
    private static String TEXT_CLASSIFICATION_TESTING = "/数据目录/testing/";

    public static void main(String[] args) throws IOException, ClassNotFoundException {
        File directory = new File("");
        try {
            TEXT_CLASSIFICATION_TESTING=directory.getCanonicalPath()+TEXT_CLASSIFICATION_TESTING;
            MODEL_FILE=directory.getCanonicalPath()+MODEL_FILE;
        } catch (IOException e) {
            e.printStackTrace();
        }
        List<List<FileModel>> listdatas=redFile("txt");
        evaluate(listdatas);
    }

    public static void evaluate(List<List<FileModel>> listdatas) throws IOException, ClassNotFoundException {
        ScoredClassifier<CharSequence> compiledClassifier = null;
        try {
            ObjectInputStream oi = new ObjectInputStream(new FileInputStream(MODEL_FILE));
            compiledClassifier = (ScoredClassifier<CharSequence>) oi.readObject();
            oi.close();
        } catch (IOException ie) {
            System.out.println("读取文件错误: 训练模型文件 " + MODEL_FILE + " 错误");
        }
        // 遍历分类目录中的文件测试分类准确度
        ConfusionMatrix confMatrix = new ConfusionMatrix(CATEGORIES);
        NumberFormat nf = NumberFormat.getInstance();
        nf.setMaximumIntegerDigits(1);
        nf.setMaximumFractionDigits(3);

        for (int i = 0; i < listdatas.size(); i++) {
            for (int j = 0; j <listdatas.get(i).size(); j++) {
                String segWords = "";
                //读取内容后分词
                List<Term> terms = HanLP.segment(listdatas.get(i).get(j).getText());
                for (Term term : terms)
                {
                    segWords += term.word + " ";
                }
                System.out.println("测试 " + CATEGORIES[i] + File.separator + listdatas.get(i).get(j).getName());
                ScoredClassification classification = compiledClassifier .classify(segWords.subSequence(0, segWords.length()));
                confMatrix.increment(CATEGORIES[i],classification.bestCategory());

                Integer b=0;

                for (int k = 0; k < CATEGORIES.length; k++) {
                    Double d= classification.score(k);
                    if(Double.isNaN(d)){
                        b++;
                    }
                }
                if(b>=CATEGORIES.length){
                    System.out.println("找不到分类，数据有问题：" +classification.category(0));
                }else{
                    System.out.println("最适合的分类: " + classification.bestCategory());
                }
                System.out.println("-------");


            }
        }


//        System.out.println("--------------------------------------------");
//        System.out.println("- 结果 ");
//        System.out.println("--------------------------------------------");
//        int[][] imatrix = confMatrix.matrix();
//        StringBuffer sb = new StringBuffer();
//        sb.append(StringTools.fillin("CATEGORY", 10, true, ' '));
//        for (int i = 0; i < CATEGORIES.length; i++)
//        {
//            sb.append(StringTools.fillin(CATEGORIES[i], 8, false, ' '));
//        };
//        System.out.println(sb.toString());
//
//        for (int i = 0; i < imatrix.length; i++) {
//            sb = new StringBuffer();
//            sb.append(StringTools.fillin(CATEGORIES[i], 10, true, ' ',10 - CATEGORIES[i].length()));
//            for (int j = 0; j < imatrix.length; j++) {
//                String out = "" + imatrix[i][j];
//                sb.append(StringTools.fillin(out, 8, false, ' ', 8 - out.length()));
//            }
//            System.out.println(sb.toString());
//        }
        System.out.println("===========");
        System.out.println("准确度: "+ nf.format(confMatrix.totalAccuracy()));
        System.out.println("总共正确数 : " + confMatrix.totalCorrect());
        System.out.println("总数：" + confMatrix.totalCount());
    }



    /**
     * 读取文件方法
     * @param type 文件后缀 .txt
     * @return
     */
    public  static List<List<FileModel>> redFile(String type){
        List<List<FileModel>> listdatas=new ArrayList<List<FileModel>>();
        for (String fl : CATEGORIES){
            List<FileModel> datas=new ArrayList<FileModel>();
            String fileurl=(TEXT_CLASSIFICATION_TESTING+fl+"/").replaceAll("\\\\","/");
            if(!FileUtil.isDirectory(fileurl)){
                String msg = "找不到训练目录" + fileurl;
                System.out.println(msg);
                throw new IllegalArgumentException(msg);
            }
            List<String> list=FileUtil.listFileNames(fileurl);
            for (int i = 0; i < list.size(); i++) {
                if("doc".equals(type.toLowerCase())||"docx".equals(type.toLowerCase())){
                    String str= readWord(fileurl+"/"+list.get(i));
                    datas.add(new FileModel(list.get(i),str,fileurl));
                }
                if("txt".equals(type.toLowerCase())){
                    String str=FileUtil.readString(fileurl+"/"+list.get(i), StandardCharsets.UTF_8 );
                    datas.add(new FileModel(list.get(i),str,fileurl));
                }
            }
            listdatas.add(datas);
        }
        return listdatas;

    }


    public static String readWord(String path) {
        String filePath = path;

        String suffixName = filePath.substring(filePath.lastIndexOf("."));//从最后一个.开始截取。截取fileName的后缀名
        try {
            File file = new File(filePath);
            FileInputStream fs = new FileInputStream(file);
            if(suffixName.equalsIgnoreCase(".doc")){//doc
                StringBuilder result2 = new StringBuilder();
                WordExtractor re = new WordExtractor(fs);
                result2.append(re.getText());//获取word中的文本内容
                re.close();
                return result2.toString();

            }else{//docx
                XWPFDocument doc = new XWPFDocument(fs);
                XWPFWordExtractor extractor = new XWPFWordExtractor(doc);
                String text = extractor.getText();//获取word中的文本内容
                extractor.close();
                fs.close();
                return text;
            }
        }catch (Exception e) {
            e.printStackTrace();
        }
        return "";

    }


}
